import torch
import matplotlib.pyplot as plt
from torch import nn

# ==================== 构建数据分布 ===================
x = torch.linspace(0, 2 * torch.pi, 100)
y = torch.sin(x) + 2
y += torch.normal(0, 0.2, y.shape)

plt.plot(torch.arange(0, len(y), 1).numpy(), y.detach().numpy(), "-", label="real")
# ==================== 设计序列模型数据 ===================
# 输入序列为5，输出序列为5，但是有输出间隔space == 1, 5, 10, 15 ...

input_size = 5
output_size = 5


def create_space_plot(space=1):
    # 计算100条数据中有多少个连续的6条数据
    steps = 100 - (input_size + output_size + space) + 1
    # 定义所选的数据大小
    data = torch.zeros((steps, input_size + output_size + space))
    # 通过顺序关系，将数据取出
    for i in range(steps):
        data[i] = y[i:i + (input_size + output_size + space)]
    inputs = data[:, :input_size]
    outputs = data[:, -output_size:]
    # ==================== 设计序列模型 ===================
    model = nn.Sequential(
        nn.Linear(input_size, 10),
        nn.Tanh(),
        nn.Linear(10, output_size)
    )
    model.train()
    # ==================== 训练 ===================
    # 损失函数
    criterion = nn.MSELoss()
    # 优化器
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    epochs = 1000
    for epoch in range(epochs):
        optimizer.zero_grad()
        predicts = model(inputs)
        loss = criterion(predicts, outputs)
        loss.backward()
        optimizer.step()

        print(f"epoch:{epoch + 1}/{epochs} -- loss:{loss.item():.4f}")
    # ==================== 序列预测 ===================
    model.eval()
    predicts = model(inputs)
    plt.plot(torch.arange(0, len(predicts), 1).numpy(), predicts[:, 0].detach().numpy(), "-", label=f"space{space}")


create_space_plot(space=1)
create_space_plot(space=15)
create_space_plot(space=20)
create_space_plot(space=35)

plt.legend()
plt.show()
